/*
池化层头文件
定义最大池化和平均池化操作
支持特征图的下采样和维度缩减
*/
#ifndef POOLINGLAYER_H
#define POOLINGLAYER_H

#include "Macros.h"
#ifdef USE_USER_DEFINED_TENSOR 
    #undef USE_EIGEN_TENSOR
    #include "Tensor.h"
#endif
#ifdef USE_EIGEN_TENSOR
    #undef USE_USER_DEFINED_TENSOR
    #include <Eigen/Core>
    #include <Eigen/Dense>
#endif
#include "Layer.h"
#include "ConvLayer.h"
#include "FullConnectedLayer.h"
#include "ActivationFunction.h"
#include "LossFunction.h"
#include <vector>
#include <memory>

#ifdef USE_EIGEN_TENSOR
// Enable Eigen parallelization
#define EIGEN_USE_THREADS
#endif

using namespace std;
#ifdef USE_EIGEN_TENSOR
using namespace Eigen;
#endif
#ifdef USE_USER_DEFINED_TENSOR
using namespace UserDefinedTensor;
#endif

class ConvLayer;
class FullConnectedLayer;

enum PoolingType {
    MAX_POOLING,
    MIN_POOLING,
    AVG_POOLING
};

void PoolingPolicy(const PoolingType& poolingType, const Tensor<double, 3>& inputFeatureMap, Tensor<double, 3>& outputFeatureMap, int widthPoolingSize, int heightPoolingSize, int stride);

class PoolingLayer : public Layer {
public:
    PoolingLayer(int layerIndex, const LayerType& layerType, int inputWidth, int inputHeight, int InputChannels, const PoolingType& poolingType, int widthPoolingSize, int heightPoolingSize, int stride=2);
    ~PoolingLayer();
    void setPrevLayer(const shared_ptr<Layer>& prevLayer) override;
    void setNextLayer(const shared_ptr<Layer>& nextLayer) override;
    const PoolingType& getPoolingType() const;
    int getWidthPoolingSize() const;
    int getHeightPoolingSize() const;
    int getStride() const;
    const Tensor<double, 3>& getInputFeatureMap() const;
    int getInputWidth() const;
    int getInputHeight() const;
    int getInputChannels() const;
    const Tensor<double, 3>& getOutputFeatureMap() const;
    const Tensor<double, 3>& getOutputFeatureMapForConv() const;
    const Tensor<double, 1>& getOutputFeatureMapForFc() const;
    int getOutputWidth() const;
    int getOutputHeight() const;
    int getOutputChannels() const;
    int getOutputWidthForConv() const;
    int getOutputHeightForConv() const;
    int getOutputChannelsForConv() const;
    int getOutputSizeForFc() const;
    const Tensor<double, 3>& getInputsGradient() const;
    const Tensor<double, 3>& getOutputsGradient() const;
    const Tensor<double, 3>& getWeightsGradient() const;
    const Tensor<double, 1>& getBiasesGradient() const;
    const Tensor<double, 3>& getInputsDelta() const;
    const Tensor<double, 3>& getOutputsDelta() const;
    const Tensor<double, 3>& getWeightsDelta() const;
    const Tensor<double, 1>& getBiasesDelta() const;
    void setPoolingType(const PoolingType& poolingType);
    void setWidthPoolingSize(int widthPoolingSize);
    void setHeightPoolingSize(int heightPoolingSize);
    void setStride(int stride);
    void setInputFeatureMap(const Tensor<double, 3>& inputFeatureMap);
    void setInputWidth(int inputWidth);
    void setInputHeight(int inputHeight);
    void setInputChannels(int inputChannels);
    void setOutputFeatureMap(const Tensor<double, 3>& outputFeatureMap);
    void setOutputFeatureMapForConv(const Tensor<double, 3>& outputFeatureMapForConv);
    void setOutputFeatureMapForFc(const Tensor<double, 1>& outputFeatureMapForFc);
    void setOutputWidth(int outputWidth);
    void setOutputHeight(int outputHeight);
    void setOutputChannels(int outputChannels);
    void setOutputWidthForConv(int outputWidthForConv);
    void setOutputHeightForConv(int outputHeightForConv);
    void setOutputChannelsForConv(int outputChannelsForConv);
    void setOutputSizeForFc(int outputSizeForFc);
    void calculateOutputDimensions(int inputWidth, int inputHeight, int inputChannels, int widthPoolingSize, int heightPoolingSize, int stride);
    void calculateOutputFeatureMap();
    void transformOutputFeatureMapToConv(Tensor<double, 3>& outputFeatureMap);
    void transformOutputFeatureMapToFc(Tensor<double, 3>& outputFeatureMap);
    void calculateGradient(Tensor<double, 3>& calcGradient, const Tensor<double, 3>& nextGradient, const Tensor<double, 3>& featureMap, const int width, const int height, const int channels);
    void calculateInputsGradient(const Tensor<double, 3>& outputFeatureMap);
    void calculateOutputsGradient(const Tensor<double, 3>& nextLayerOutputsGradient);
    void calculateInputsDelta(const Tensor<double, 3>& outputsDelta);
    void calculateOutputsDelta(int learningRate, int momentum);
    void forward(const Tensor<double, 3>& inputFeatureMap);
    void backward(const Tensor<double, 3>& nextLayerOutputsGradient);
    void print() const;
    
private:
    PoolingType poolingType;
    int widthPoolingSize;
    int heightPoolingSize;
    int stride;
    Tensor<double, 3> inputFeatureMap;
    int inputWidth;
    int inputHeight;
    int inputChannels;
    Tensor<double, 3> outputFeatureMap;
    Tensor<double, 3> outputFeatureMapForConv;
    Tensor<double, 1> outputFeatureMapForFc;
    int outputWidth;
    int outputHeight;
    int outputChannels;
    int outputWidthForConv;
    int outputHeightForConv;
    int outputChannelsForConv;
    int outputSizeForFc;
    LayerType nextLayerType;
    shared_ptr<ConvLayer> prevConvLayer;
    shared_ptr<ConvLayer> nextConvLayer;
    shared_ptr<FullConnectedLayer> nextFullConnectedLayer;
    Tensor<double, 3> inputsGradient;
    Tensor<double, 3> outputsGradient;
    Tensor<double, 3> weightsGradient;
    Tensor<double, 1> biasesGradient;
    Tensor<double, 3> inputsDelta;
    Tensor<double, 3> outputsDelta;
    Tensor<double, 3> weightsDelta;
    Tensor<double, 1> biasesDelta;
    Tensor<double, 3> prevBatchInputsDelta;
    Tensor<double, 3> prevBatchOutputsDelta;
    Tensor<double, 3> prevBatchWeightsDelta;
    Tensor<double, 1> prevBatchBiasesDelta;
};

#endif // POOLINGLAYER_H